import pandas as pd
import os
import numpy as np  # 引入numpy用于处理可能的无穷大值
import plotly.graph_objects as go
from plotly.subplots import make_subplots


# --- 数据加载函数 (与您的模板完全相同) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从指定的CSV文件中加载时间范围内的数据。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]
    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                # 只读取我们需要的列以优化内存
                df = pd.read_csv(file_path, usecols=['open_time', 'open', 'close', 'volume'])
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


# --- 数据处理与指标计算函数 ---
def process_and_calculate_metrics(df: pd.DataFrame, start_time_str: str, end_time_str: str) -> pd.DataFrame:
    """
    对数据进行预处理，并计算核心分析指标。
    """
    if df.empty:
        return df

    # 基础处理
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'close', 'volume']
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)
    df_filtered = df.loc[start_time_str:end_time_str].copy()  # 使用 .copy() 避免 SettingWithCopyWarning

    if df_filtered.empty:
        print("在指定时间范围内没有数据。")
        return pd.DataFrame()

    # --- 核心指标计算 ---
    # 计算价格变化的绝对值比例
    # 使用 np.divide 来处理分母为0的情况，虽然在价格数据中极罕见
    df_filtered['price_change_ratio'] = np.divide(
        (df_filtered['close'] - df_filtered['open']).abs(),
        df_filtered['close']
    )

    # 清理计算中可能产生的无穷大或NaN值
    df_filtered.replace([np.inf, -np.inf], np.nan, inplace=True)
    df_filtered.dropna(subset=['price_change_ratio', 'volume'], inplace=True)

    print(f"数据处理和指标计算完成。共 {len(df_filtered)} 条有效K线。")
    return df_filtered


# --- 直方图绘制函数 ---
def plot_distributions(df: pd.DataFrame, chart_title: str):
    """
    绘制 Volume 和 Price Change Ratio 的分布直方图。
    """
    if df.empty:
        print("数据为空，无法绘制直方图。")
        return

    fig = make_subplots(
        rows=2, cols=1,
        subplot_titles=(
            "<b>交易量 (Volume) 分布</b>",
            "<b>价格变化率 (abs(close-open)/close) 分布</b>"
        ),
        vertical_spacing=0.15
    )

    # 1. 绘制 Volume 直方图
    fig.add_trace(go.Histogram(
        x=df['volume'],
        name='Volume',
        marker_color='royalblue',
        opacity=0.75
    ), row=1, col=1)

    # 为了更好的可视化，我们可能需要对价格变化率的极端离群值进行裁剪
    # 这里我们只显示99.5%分位数以内的数据，以观察主体分布
    price_change_viz_limit = df['price_change_ratio'].quantile(0.995)
    df_viz = df[df['price_change_ratio'] < price_change_viz_limit]

    # 2. 绘制 Price Change Ratio 直方图
    fig.add_trace(go.Histogram(
        x=df_viz['price_change_ratio'],
        name='Price Change Ratio',
        marker_color='darkorange',
        opacity=0.75
    ), row=2, col=1)

    # --- 在图上添加统计注释线 ---
    # Volume 统计线
    vol_mean = df['volume'].mean()
    vol_median = df['volume'].median()
    fig.add_vline(x=vol_mean, line_width=2, line_dash="dash", line_color="red",
                  annotation_text=f"均值: {vol_mean:,.0f}", annotation_position="top right",
                  row=1, col=1)
    fig.add_vline(x=vol_median, line_width=2, line_dash="dash", line_color="green",
                  annotation_text=f"中位数: {vol_median:,.0f}", annotation_position="bottom right",
                  row=1, col=1)

    # Price Change Ratio 统计线
    price_change_mean = df['price_change_ratio'].mean()
    price_change_median = df['price_change_ratio'].median()
    fig.add_vline(x=price_change_mean, line_width=2, line_dash="dash", line_color="red",
                  annotation_text=f"均值: {price_change_mean:.5f}", annotation_position="top right",
                  row=2, col=1)
    fig.add_vline(x=price_change_median, line_width=2, line_dash="dash", line_color="green",
                  annotation_text=f"中位数: {price_change_median:.5f}", annotation_position="bottom right",
                  row=2, col=1)

    # 更新整体布局
    start_str = df.index[0].strftime('%Y-%m-%d %H:%M')
    end_str = df.index[-1].strftime('%Y-%m-%d %H:%M')
    full_title = f'{chart_title} 指标分布统计 ({start_str} to {end_str})'

    fig.update_layout(
        title_text=full_title,
        showlegend=False,
        height=800,
        bargap=0.05
    )
    fig.update_yaxes(title_text='频数 (Frequency)', row=1, col=1)
    fig.update_yaxes(title_text='频数 (Frequency)', row=2, col=1)
    fig.update_xaxes(title_text='Volume', row=1, col=1)
    fig.update_xaxes(title_text='abs(close-open)/close (已裁剪99.5%以上离群值)', row=2, col=1)

    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "ETHUSDT"
    TIMEFRAME = "5m"
    # 建议选择一个较长的时间跨度（如数月）以获得有统计意义的分布
    START_TIME = "2025-01-01 00:00:00"
    END_TIME = "2025-05-09 00:00:00"
    # ====================================================

    print("--- 数据分布分析程序启动 ---")
    raw_data_df = load_data_in_range(START_TIME, END_TIME, BASE_DATA_PATH, SYMBOL, TIMEFRAME)
    final_df = process_and_calculate_metrics(raw_data_df, START_TIME, END_TIME)

    if not final_df.empty:
        chart_title_str = f"{SYMBOL} - {TIMEFRAME}"
        plot_distributions(final_df, chart_title=chart_title_str)
        print("--- 直方图生成完毕 ---")
    else:
        print("\n未能加载或处理任何数据，程序退出。请检查您的配置和文件路径。")